#!/usr/bin/env python
# -*- coding:utf-8 _*-
"""
@file: decoder_surgeon.py
@time: 2023/07/20
@author: xingwg
@contact: xwg031459@163.com
@software: PyCharm
"""
from copy import deepcopy
import numpy as np
import onnx
import onnx_graphsurgeon as gs

src_onnx = "onnx/decoder_opt.onnx"
dst_onnx = "decoder_opt.onnx"

graph = gs.import_onnx(onnx.shape_inference.infer_shapes(onnx.load(src_onnx)))

cnt = 0
for node in graph.nodes:
    if node.op == "Reshape" and node.o().op == "InstanceNormalization" and node.o().o().op == "Reshape" \
            and node.o().o().o().op == "Mul" and node.o().o().o().o().op == "Add":

        last_node = node.o().o().o().o()

        instance_norm = node.o()
        instance_norm_scale = instance_norm.inputs[1]
        instance_norm_bias = instance_norm.inputs[2]
        epsilon = instance_norm.attrs["epsilon"]
        mul_node = node.o().o().o()
        add_node = node.o().o().o().o()

        scale = np.ascontiguousarray(np.array(deepcopy(instance_norm_scale.values.tolist()), dtype=np.float32))
        bias = np.ascontiguousarray(np.array(deepcopy(instance_norm_bias.values.tolist()), dtype=np.float32))
        gamma = np.ascontiguousarray(np.array(deepcopy(mul_node.inputs[1].values.tolist()), dtype=np.float32))
        beta = np.ascontiguousarray(np.array(deepcopy(add_node.inputs[1].values.tolist()), dtype=np.float32))

        with_swish = True if node.o().o().o().o().o().op == "Sigmoid" and node.o().o().o().o().o().o().op == "Mul" else False
        if with_swish:
            last_node = node.o().o().o().o().o().o()

        constant_gamma = gs.Constant("gamma_{}".format(cnt), gamma.reshape(-1))
        constant_beta = gs.Constant("beta_{}".format(cnt), beta.reshape(-1))
        # constant_scale = gs.Constant("scale_{}".format(cnt), scale.reshape(-1))
        # constant_bias = gs.Constant("bias_{}".format(cnt), bias.reshape(-1))
        x = node.inputs[0]
        group_norm_v = gs.Variable("group_norm_{}".format(cnt), np.dtype(np.float32), x.shape)
        group_norm = gs.Node("GroupNorm", "GroupNorm_{}".format(cnt),
                             # attrs={"epsilon": epsilon, "num_groups": instance_norm_scale.values.size, "swish": with_swish},
                             attrs={"epsilon": epsilon, "bSwish": with_swish},
                             inputs=[x, constant_gamma, constant_beta],
                             outputs=[group_norm_v])
        cnt += 1
        for n in graph.nodes:
            if last_node.outputs[0] in n.inputs:
                index = n.inputs.index(last_node.outputs[0])
                n.inputs[index] = group_norm.outputs[0]
        last_node.outputs = []
        graph.nodes.append(group_norm)

graph.cleanup()
onnx.save(gs.export_onnx(graph), dst_onnx)
